import tkinter as tk
import random
from math import sqrt
from collections import defaultdict

# 定义蚂蚁类，用于表示单个蚂蚁的属性和状态
class Ant:
    def __init__(self, aid, home, color):
        # 蚂蚁的唯一标识符
        self.id = aid
        # 蚂蚁当前所在的位置，初始位置为蚁穴位置
        self.pos = home
        # 蚂蚁走过的路径，存储其经过的位置
        self.path = []
        # 标记蚂蚁是否携带食物，初始为 False
        self.has_food = False
        # 蚂蚁的颜色，用于在界面上区分不同的蚂蚁
        self.color = color
        # 蚂蚁在画布上对应的图形对象，初始为 None
        self.obj = None

# 定义蚁群类，用于管理整个蚁群模拟环境
class AntColony:
    def __init__(self, master):
        # 主窗口对象，用于创建界面
        self.master = master
        # 设置主窗口的标题
        self.master.title("改进蚁群算法")
        
        # 参数配置
        # 地图的尺寸，即地图的边长（假设地图为正方形）
        self.size = 15
        # 显示倍率，用于控制地图在界面上的显示大小
        self.scale = 50
        # 蚁穴的位置，用坐标 (x, y) 表示
        self.home = (7, 8)
        # 食物的位置，用坐标 (x, y) 表示
        self.food = (1, 1)
        # 蚂蚁的数量
        self.num_ants = 30
        # 信息素的挥发率，用于模拟信息素随时间的衰减
        self.evaporation = 0.1
        # 信息素权重，在蚂蚁选择移动方向时，影响信息素浓度的重要程度
        self.alpha = 1
        # 启发式权重，在蚂蚁选择移动方向时，影响到食物距离的重要程度
        self.beta = 2
        
        # 初始化界面
        # 创建一个画布对象，用于绘制地图、蚂蚁、食物和蚁穴等元素
        self.canvas = tk.Canvas(master, width=self.size * self.scale, 
                               height=self.size * self.scale, bg='white')
        # 将画布添加到主窗口中，并设置垂直方向的间距
        self.canvas.pack(pady=10)
        # 创建一个重置按钮，点击该按钮将调用 reset 方法重置模拟环境
        tk.Button(master, text="重置", command=self.reset).pack()
        
        # 调用重置方法，初始化模拟环境
        self.reset()
        self.update()
    
    def reset(self):
        """重置模拟环境"""
        # 删除画布上的所有元素
        self.canvas.delete("all")
        # 初始化信息素字典，使用 defaultdict 确保访问不存在的键时返回默认值 0.0
        self.pheromones = defaultdict(float)
        # 创建指定数量的蚂蚁对象，并存储在列表中
        self.ants = [Ant(i, self.home, self.random_color()) 
                    for i in range(self.num_ants)]
        
        # 绘制食物和蚁穴
        # 在画布上绘制表示食物的圆形点，颜色为橙色
        self.draw_point(self.food, 'orange')
        # 在画布上绘制表示蚁穴的圆形点，颜色为棕色
        self.draw_point(self.home, 'brown')
        
        # 初始化蚂蚁显示
        # 遍历所有蚂蚁，在画布上绘制每个蚂蚁的圆形点，并将图形对象赋值给蚂蚁的 obj 属性
        for ant in self.ants:
            ant.obj = self.draw_point(ant.pos, ant.color)
        
        # 调用更新方法，开始模拟循环
        

    def draw_point(self, pos, color, size=0.5):
        """绘制圆形点"""
        # 解包位置坐标
        x, y = pos
        # 在画布上创建一个圆形，根据位置和显示倍率计算坐标，并填充指定颜色
        return self.canvas.create_oval(
            (x - size) * self.scale, (y - size) * self.scale,
            (x + size) * self.scale, (y + size) * self.scale,
            fill=color, outline=''
        )

    def random_color(self):
        """生成随机颜色"""
        # 生成三个随机整数，范围在 50 到 205 之间，用于表示 RGB 颜色值
        # 然后将其转换为十六进制字符串，组合成一个有效的 HTML 颜色代码
        return f'#{random.randint(50, 205):02x}{random.randint(50, 205):02x}{random.randint(50, 205):02x}'

    def update(self):
        """主更新循环"""
        # 信息素挥发
        # 遍历信息素字典中的所有键
        for k in list(self.pheromones):
            # 根据挥发率更新信息素浓度
            self.pheromones[k] *= (1 - self.evaporation)
            # 如果信息素浓度低于 0.1，则从字典中删除该键
            if self.pheromones[k] < 0.1:
                del self.pheromones[k]
        
        # 移动蚂蚁
        # 遍历所有蚂蚁
        for ant in self.ants:
            # 调用 move_ant 方法移动当前蚂蚁
            self.move_ant(ant)
            # 更新画布上蚂蚁图形对象的位置，使其与蚂蚁的当前位置对应
            self.canvas.coords(ant.obj,
                (ant.pos[0] - 0.4) * self.scale, (ant.pos[1] - 0.4) * self.scale,
                (ant.pos[0] + 0.4) * self.scale, (ant.pos[1] + 0.4) * self.scale
            )
        
        # 每隔 500 毫秒调用一次 update 方法，实现循环更新
        self.master.after(1000, self.update)

    def move_ant(self, ant):
        """蚂蚁移动逻辑"""
        if ant.has_food:
            # 带食物返回蚁穴
            # 如果蚂蚁的路径列表不为空
            if ant.path:
                # 从路径列表中取出最后一个位置作为新的位置
                ant.pos = ant.path.pop()
                # 在该位置增加信息素浓度
                self.pheromones[f"{ant.pos[0]}_{ant.pos[1]}"] += 1
            else:
                # 如果路径列表为空，说明蚂蚁已回到蚁穴，将 has_food 标记设为 False
                ant.has_food = False
        else:
            # 寻找食物
            # 调用 choose_move 方法选择下一个移动位置
            next_pos = self.choose_move(ant)
            # 将当前位置添加到路径列表中
            ant.path.append(ant.pos)
            # 更新蚂蚁的当前位置
            ant.pos = next_pos
            
            # 发现食物
            # 如果蚂蚁移动到了食物的位置
            if ant.pos == self.food:
                # 将 has_food 标记设为 True，表示蚂蚁携带了食物
                ant.has_food = True
                # 反转路径列表，以便蚂蚁能沿着原路返回蚁穴
                ant.path = ant.path[::-1]

    def choose_move(self, ant):
        """选择移动方向"""
        # 获取蚂蚁当前的位置坐标
        x, y = ant.pos
        # 计算蚂蚁周围四个相邻位置的坐标，确保不超出地图边界
        neighbors = [
            (max(x - 1, 0), y),   # 左
            (min(x + 1, self.size - 1), y),  # 右
            (x, max(y - 1, 0)),   # 上
            (x, min(y + 1, self.size - 1)),  # 下
        ]
        
        # 计算概率
        probabilities = []
        # 遍历所有相邻位置
        for nx, ny in neighbors:
            # 生成该位置的键，用于在信息素字典中查找信息素浓度
            key = f"{nx}_{ny}"
            # 计算该位置到食物的距离的倒数，作为启发式信息
            distance = 1 / (1 + sqrt((nx - self.food[0])**2 + (ny - self.food[1])**2))
            # 获取该位置的信息素浓度，如果不存在则使用默认值 1e-4
            pheromone = self.pheromones.get(key, 1e-4)
            # 根据信息素权重和启发式权重计算选择该位置的概率
            prob = (pheromone**self.alpha) * (distance**self.beta)
            probabilities.append(prob)
        
        # 概率归一化
        # 计算所有概率的总和
        total = sum(probabilities)
        # 如果总和为 0，说明所有相邻位置的概率都为 0，随机选择一个相邻位置
        if total == 0:
            return random.choice(neighbors)
        # 将所有概率除以总和，得到归一化后的概率
        probabilities = [p / total for p in probabilities]
        
        # 根据归一化后的概率，随机选择一个相邻位置作为下一个移动位置
        return random.choices(neighbors, weights=probabilities, k=1)[0]

if __name__ == "__main__":
    # 创建主窗口对象
    root = tk.Tk()
    # 创建蚁群模拟对象
    simulation = AntColony(root)
    # 进入主事件循环，使窗口保持显示并响应用户操作
    root.mainloop()